import os
import numpy as np
import cv2

def img2vector(imgfilename):
    # 以 灰度格式读入图片，返回一维向量
    img = cv2.imread(imgfilename,cv2.IMREAD_GRAYSCALE);
    rows , columns = img.shape
    vec = img.reshape(rows * columns)
    return vec

training_dir = "../training"
sub_dir_and_files = os.listdir(training_dir)

sub_dirs = []
# 如果是目录
for x in sub_dir_and_files:
    if os.path.isdir(training_dir + "/" + x):
        sub_dirs.append(x)

# 计算装备图像总数
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(training_dir + "/" + subdir))


# 初始化训练图像数据矩阵 (N 行，128*128列) 和装备向量（长度为N）
training_img_matrix = np.zeros((N,128*128)) # 每个图像一行数据
training_equipment_vector = [''] * N

i = 0 # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(training_dir + "/" + subdir)
    for image in image_files:
        # 将图像转换为向量
        v = img2vector(training_dir + "/" + subdir + "/" + image)
        training_img_matrix[i] = v
        training_equipment_vector[i] = subdir
        i += 1

print(training_img_matrix)
print(training_equipment_vector)




